//
// Created by 高森森 on 2022/2/13.
//

#ifndef LEETCODE_SOLUTION40_H
#define LEETCODE_SOLUTION40_H
#include<iostream>
#include <vector>
#include <numeric>

using namespace std;
class UnionFind2{
public:
    vector<int>father;//父亲节点
    int n;//节点个数
    vector<int>size;//连通分量数
public:
    //查询父结点
    UnionFind2(int _n):n(_n),father(_n),size(_n,1){
        //iota函数，令father[i]=i;
        iota(father.begin(),father.end(),0);
    }
    int find(int x){
        if(x==father[x])
            return x;
        else{
            father[x]=find(father[x]);
            return father[x];
        }
    }

    //并查集合并
    bool merge(int i,int j){
        int x=find(i);
        int y=find(j);
        father[x]=y;
        size[y]+=size[x];
        return true;
    }
    bool isConnected(int x,int y){
        x=find(x);
        y=find(y);
        return x==y;
    }
};

class Solution40 {
public:
    int minMalwareSpread(vector<vector<int>>& graph, vector<int>& initial) {
        int n=graph.size();
        vector<int>colors(n,-1);
        vector<vector<int>>g(n);
        for(int i=0;i<n;i++){
           for(int j=0;j<n;j++){
               if(graph[i][j]==1&&i!=j){
                   g[i].push_back(j);
               }
           }
        }
        for(int i=0;i<n;i++)
        {
            if(colors[i]==-1)
                dfs(i,i,colors,g);
        }
        //计算每个颜色的连通分量的大小
        vector<int>size(n,0);
        for(int color:colors){
            size[color]++;
        }
        //找到initial中连通分量为1的颜色节点
        //计算节点中存在的所有颜色，找到颜色唯一的点
        vector<int>count(n);
        for(int node:initial){
            count[colors[node]]++;
        }
        int ans=INT_MAX;
        for(int node:initial){
            int c=colors[node];
            if(count[c]==1){
                if(ans==INT_MAX)
                    ans=node;
                else if(size[c]>size[colors[ans]]){
                    ans=node;
                }
                else if(size[c]==size[colors[ans]]&&node<ans){
                    ans=node;
                }
            }
        }
        if(ans==INT_MAX){
            for(int node:initial){
                ans=min(node,ans);
            }
        }
        return ans;

    }
    //dfs染色
    void dfs(int node,int color,vector<int>&colors,vector<vector<int>>graph){
        colors[node]=color;
        for(auto neigh:graph[node]){
            if(colors[neigh]==-1)
            dfs(neigh,color,colors,graph);
        }
    }
    int minMalwareSpread2(vector<vector<int>>& graph, vector<int>& initial) {
        int n = graph.size();
        UnionFind2 unionFind2(n);
        for (int i = 0; i < n; i++)
            for (int j = i + 1; j < n; j++) {
                if (graph[i][j] == 1) {
                    unionFind2.merge(i, j);
                }
            }
        vector<int> count(n);
        for (int node:initial) {
            count[unionFind2.find(node)]++;
        }
        int ans = INT_MAX;
        for (int node:initial) {
            int f = unionFind2.find(node);
            if (count[f] == 1) {
                if (ans == INT_MAX)
                    ans = node;
                else if (unionFind2.size[f] > unionFind2.size[unionFind2.find(ans)]) {
                    ans = node;
                } else if (unionFind2.size[f] == unionFind2.size[unionFind2.find(ans)] && node < ans) {
                    ans = node;
                }
            }
        }
        if(ans==INT_MAX){
            for(int node:initial){
                ans=min(node,ans);
            }
        }
        return ans;
    }
};


#endif //LEETCODE_SOLUTION40_H
